import os
import time
import xml.etree.ElementTree as ET
import click
import pandas as pd
from thefuzz import fuzz
from multiprocessing import Pool
from tqdm import tqdm, trange
from InquirerPy import inquirer


class BioConvert:
    def generate_cache(self, database, cache):
        t = time.time()
        ns = {'uniprot': 'http://uniprot.org/uniprot'}
        print("XML Database Loading...\nIt may take a few minutes, depending on the performance of your computer.")
        root = ET.parse(database).getroot()
        genes = list()
        proteins = list()
        entries = root.findall('uniprot:entry', ns)
        for entry in entries:
            gene = entry.find('uniprot:gene', ns)
            if gene is None:
                continue

            protein = entry.find('uniprot:protein', ns)
            if protein is None:
                continue

            genes.append(gene.find('uniprot:name', ns).text)
            protein_name_temp = list()
            protein_name_group = list()
            protein_name_group.append(protein)
            for i in protein.findall('uniprot:domain', ns):
                protein_name_group.append(i)
            for i in protein.findall('uniprot:component', ns):
                protein_name_group.append(i)

            for i, v in enumerate(protein_name_group):
                recommended_name = v.find('uniprot:recommendedName', ns)
                if recommended_name is not None:
                    full_name = recommended_name.find('uniprot:fullName', ns)
                    if full_name is not None:
                        protein_name_temp.append(full_name.text)
                    short_name = recommended_name.find('uniprot:shortName', ns)
                    if short_name is not None:
                        protein_name_temp.append(short_name.text)
                for alternative_name in v.findall('uniprot:alternativeName', ns):
                    full_name = alternative_name.find('uniprot:fullName', ns)
                    if full_name is not None:
                        protein_name_temp.append(full_name.text)
                    short_name = alternative_name.find('uniprot:shortName', ns)
                    if short_name is not None:
                        protein_name_temp.append(short_name.text)
            proteins.append(protein_name_temp)
        print("XML Database Loading Completed. It took ",
              time.time() - t, "seconds.")

        df = pd.DataFrame()
        df['Gene'] = genes
        df['Proteins'] = proteins
        df.to_excel(cache, index=False)

    def load_cache(self, cache):
        self.cache_df = pd.read_excel(cache, usecols=['Gene', 'Proteins'], dtype=str)
        self.db_proteins = list()
        for row in self.cache_df.itertuples():
            self.db_proteins.append(eval(row.Proteins))

    def protein_to_gene_preprocess(self, protein):
        gene = ""
        for i, proteins in enumerate(self.db_proteins):
            for j, v in enumerate(proteins):
                score = fuzz.token_sort_ratio(protein, v, full_process=False)
                if score == 100:
                    if len(gene) != 0:
                        gene += "/"
                    gene += self.cache_df.loc[i, "Gene"]
                    break
        return gene

    def convert(self, uniprot, tcmsp, output, processes):
        if os.path.exists(output):
            proceed = inquirer.confirm(
                message="输出文件已经存在，继续将覆盖它！是否继续？").execute()
            if not proceed:
                return

        cache = uniprot + '.xlsx'
        if not os.path.exists(cache):
            self.generate_cache(uniprot, cache)
        self.load_cache(cache)
        tcmsp_df = pd.read_excel(
            tcmsp, usecols=['Target name'], keep_default_na=False, dtype=str)
        tcmsp_unique_df = tcmsp_df.copy().drop_duplicates()
        with Pool(processes=processes) as p:
            tcmsp_unique_df['Gene name'] = tqdm(p.imap(
                self.protein_to_gene_preprocess, tcmsp_unique_df['Target name']), total=tcmsp_unique_df['Target name'].count(), desc='去重匹配')
        for i in trange(len(tcmsp_df), desc="最终匹配"):
            protein = tcmsp_df.loc[i, 'Target name']
            gene = tcmsp_unique_df.loc[tcmsp_unique_df["Target name"]
                                       == protein].values[0][1]
            tcmsp_df.at[i, "Gene name"] = gene

        tcmsp_df.to_excel(output, index=False)


@click.command()
@click.option('--uniprot', required=True, help="XML from uniprot website.")
@click.option('--tcmsp', required=True, help="An Excel with column \"Target names\", manual collect from tcmsp website.")
@click.option('--output', default="output.xlsx", help="An Excel with convert result.")
@click.option('--processes', type=int, default=None, help="processes is the number of worker processes to use. If processes is None then the number returned by os.cpu_count() is used.")
def cli_convert(uniprot, tcmsp, output, processes):
    """
    python protein2gene.py --uniprot uniprot-file --tcmsp tcmsp-file [--output output-file] [--processes 4]

    uniprot-file: *.xml
    tcmsp-file: *.xlsx
    output-file: *.xlsx
    """
    bio_convert = BioConvert()
    bio_convert.convert(uniprot, tcmsp, output, processes)


if __name__ == "__main__":
    cli_convert()
